'''示例：

export FAKE_PHOTOS_RWKV_WORKER="$(pwd)/rwkv_worker.py"
export FAKE_PHOTOS_TMP_PROMPT="b.txt"
export FAKE_PHOTOS_SD_EXE="/var/fake-photos/stable-diffusion.cpp/bin/sd"
python rwkv_client.py -o a.png -p "霞光照射在高楼上" --steps 1 -s 42 -W 512 -H 512 -m "/var/fake-photos/sd_xl_turbo_1.0_fp16.safetensors" --vae "/var/fake-photos/sdxl.vae.safetensors" --cfg-scale 1 -t 6 -v

'''

import subprocess
import argparse
import os
import sys

# 文件 rwkv_worker.py 的路径
FAKE_PHOTOS_RWKV_WORKER = os.environ.get('FAKE_PHOTOS_RWKV_WORKER')
if FAKE_PHOTOS_RWKV_WORKER is None:
    print('FAKE_PHOTOS_RWKV_WORKER is not set!', file=sys.stderr)
    exit(1)

# 临时文件 prompt.txt 的路径
FAKE_PHOTOS_TMP_PROMPT = os.environ.get('FAKE_PHOTOS_TMP_PROMPT')
if FAKE_PHOTOS_TMP_PROMPT is None:
    print('FAKE_PHOTOS_TMP_PROMPT is not set!', file=sys.stderr)
    exit(1)

# sd可执行程序路径
FAKE_PHOTOS_SD_EXE = os.environ.get('FAKE_PHOTOS_SD_EXE')
if FAKE_PHOTOS_SD_EXE is None:
    print('FAKE_PHOTOS_SD_EXE is not set!', file=sys.stderr)
    exit(1)

parser = argparse.ArgumentParser(description='Image generation with RWKV')
parser.add_argument('-o', help='output image')
parser.add_argument('-p', help='prompt')
parser.add_argument('--steps', help='sample steps')
parser.add_argument('-s', help='seed, if < 0 use random seed')
parser.add_argument('-W', help='image width')
parser.add_argument('-H', help='image height')
# 是否不使用rwkv直接传提示词给sd 可选
parser.add_argument('--norwkv', help='do not use rwkv')

args, ext = parser.parse_known_args()
print(f"output image: {args.o}")
print(f"prompt: {args.p}")
print(f"steps: {args.steps}")
print(f"width: {args.W}")
print(f"height: {args.H}")

def model(inp, seed):
    subprocess.run([
        'python', FAKE_PHOTOS_RWKV_WORKER,
        '-o', FAKE_PHOTOS_TMP_PROMPT,
        '-p', inp,
        '-s', seed
    ])
    modelOutput = ''
    with open(FAKE_PHOTOS_TMP_PROMPT, 'r') as f:
        modelOutput = f.read()
    return modelOutput

modelOutput = args.p
if args.norwkv == '0':
    modelOutput = model(args.p, args.s)

print(modelOutput)

command = [FAKE_PHOTOS_SD_EXE] + ext + [
    '-o', args.o,
    '-p', modelOutput,
    '--steps', args.steps,
    '-s', args.s,
    '-W', args.W,
    '-H', args.H
]

print('Command: ' + ' '.join(command))

# 检查退出状态，如果不是0那么输出错误到标准输出
if subprocess.run(command).returncode != 0:
    print('Error: sd exited with non-zero status, command was' + ' '.join(command), file=sys.stderr)
    exit(1)
